#!/usr/bin/env python3
"""验证环境配置是否正确"""

import sys
import platform

def print_section(title):
    """打印分节标题"""
    print(f"\n{'='*60}")
    print(f"  {title}")
    print('='*60)

def check_import(module_name, package_name=None, version_attr='__version__'):
    """检查模块是否可导入并显示版本"""
    if package_name is None:
        package_name = module_name
    
    try:
        module = __import__(module_name)
        version = getattr(module, version_attr, 'unknown')
        print(f"✓ {package_name:20s} - version: {version}")
        return True
    except ImportError as e:
        print(f"✗ {package_name:20s} - 导入失败: {str(e)[:50]}")
        return False
    except Exception as e:
        print(f"⚠ {package_name:20s} - 警告: {str(e)[:50]}")
        return True

def main():
    """主验证流程"""
    
    # 系统信息
    print_section("系统信息")
    print(f"操作系统: {platform.system()} {platform.release()}")
    print(f"Python 版本: {sys.version}")
    print(f"Python 路径: {sys.executable}")
    
    # 核心依赖
    print_section("核心依赖包")
    core_packages = [
        ("numpy", "NumPy"),
        ("scipy", "SciPy"),
        ("torch", "PyTorch"),
        ("torchvision", "TorchVision"),
        ("jax", "JAX"),
        ("jaxlib", "JAXLib"),
        ("casadi", "CasADi"),
        ("matplotlib", "Matplotlib"),
    ]
    
    core_success = sum(check_import(m, n) for m, n in core_packages)
    
    # 优化求解器
    print_section("优化求解器")
    solver_packages = [
        ("julia", "PyJulia"),
    ]
    
    solver_success = sum(check_import(m, n) for m, n in solver_packages)
    
    # 实验工具
    print_section("实验和可视化工具")
    exp_packages = [
        ("wandb", "Weights & Biases"),
        ("tensorboard", "TensorBoard"),
        ("tqdm", "TQDM"),
    ]
    
    exp_success = sum(check_import(m, n) for m, n in exp_packages)
    
    # 其他依赖
    print_section("其他依赖")
    other_packages = [
        ("pydantic", "Pydantic"),
        ("yaml", "PyYAML"),
        ("requests", "Requests"),
        ("click", "Click"),
    ]
    
    other_success = sum(check_import(m, n) for m, n in other_packages)
    
    # 项目模块
    print_section("项目模块")
    project_success = 0
    
    try:
        sys.path.insert(0, '.')
        from src.vae.vae_model import VAE
        print("✓ src.vae.vae_model     - VAE 模型")
        project_success += 1
    except Exception as e:
        print(f"✗ src.vae.vae_model     - {str(e)[:50]}")
    
    try:
        from src.vae.encoder import Encoder
        print("✓ src.vae.encoder       - 编码器")
        project_success += 1
    except Exception as e:
        print(f"✗ src.vae.encoder       - {str(e)[:50]}")
    
    try:
        from src.vae.decoder import Decoder
        print("✓ src.vae.decoder       - 解码器")
        project_success += 1
    except Exception as e:
        print(f"✗ src.vae.decoder       - {str(e)[:50]}")
    
    try:
        from src.mcp.mcp_solver import MCPSolver
        print("✓ src.mcp.mcp_solver    - MCP 求解器")
        project_success += 1
    except Exception as e:
        print(f"✗ src.mcp.mcp_solver    - {str(e)[:50]}")
    
    try:
        from games.drone_game import DroneGame
        print("✓ games.drone_game      - 无人机博弈")
        project_success += 1
    except Exception as e:
        print(f"✗ games.drone_game      - {str(e)[:50]}")
    
    # GPU 支持
    print_section("GPU 支持")
    import torch
    if torch.cuda.is_available():
        print(f"✓ CUDA 可用")
        print(f"  设备名称: {torch.cuda.get_device_name(0)}")
        print(f"  设备数量: {torch.cuda.device_count()}")
        print(f"  CUDA 版本: {torch.version.cuda}")
    else:
        print("✗ CUDA 不可用 (将使用 CPU 进行计算)")
    
    # JAX GPU 支持
    try:
        import jax
        devices = jax.devices()
        print(f"✓ JAX 设备: {devices}")
    except Exception as e:
        print(f"⚠ JAX 设备检查失败: {e}")
    
    # Julia 配置
    print_section("Julia 配置")
    try:
        import julia
        import os
        julia_bindir = os.environ.get('JULIA_BINDIR', 'Not set')
        print(f"  JULIA_BINDIR: {julia_bindir}")
        
        # 尝试初始化 Julia
        try:
            from julia import Main
            Main.eval('println("Julia 运行正常")')
            print("✓ Julia 初始化成功")
        except Exception as e:
            print(f"✗ Julia 初始化失败: {e}")
            print("  提示: 请运行 python -c 'import julia; julia.install()'")
    except ImportError:
        print("✗ PyJulia 未安装")
    
    # 总结
    print_section("验证总结")
    total_core = len(core_packages)
    total_solver = len(solver_packages)
    total_exp = len(exp_packages)
    total_other = len(other_packages)
    total_project = 5
    
    print(f"核心依赖:   {core_success}/{total_core}")
    print(f"求解器:     {solver_success}/{total_solver}")
    print(f"实验工具:   {exp_success}/{total_exp}")
    print(f"其他依赖:   {other_success}/{total_other}")
    print(f"项目模块:   {project_success}/{total_project}")
    
    total_success = core_success + solver_success + exp_success + other_success + project_success
    total_packages = total_core + total_solver + total_exp + total_other + total_project
    
    percentage = (total_success / total_packages) * 100
    print(f"\n总体完成度: {total_success}/{total_packages} ({percentage:.1f}%)")
    
    if percentage == 100:
        print("\n🎉 环境配置完美！所有依赖都已正确安装。")
        return 0
    elif percentage >= 80:
        print("\n✓ 环境基本可用，但有一些可选依赖缺失。")
        return 0
    else:
        print("\n⚠ 环境配置不完整，请检查缺失的依赖。")
        return 1

if __name__ == "__main__":
    sys.exit(main())
